大家好，昨天我們把圖片抓下來之後也標記完了(我個人是用了10000張圖片)，接下來就是丟進模型訓練啦！這邊小弟採用pytorch的框架，我們首先要寫一個讀取資料集的程式。

• 圖片整理

我們要將抓下來的10000張圖片，分成train、test及validation，我是切7、2、1。
test可以邊訓練邊觀察有沒有練起來及有沒有過擬和，而validation則是讓你測試模型準確度及泛化程度好不好。

• 建立dataset

首先要import我們需要的套件。

1. 我們在讀取圖片資料夾的時候會用到os.listdir將圖片名稱產生成一個list。
2. PIL則是用在開啟圖片及轉換RGB或者調整圖片大小時可用到。
3. torch.utils.data則是助於我們建立Dataset以供訓練模型時讀取。
4. pandas則是我在導入圖片的label時會需要用到。

import os
from PIL import Image
import torch
from torch.utils.data import Dataset
import pandas as pd

首先我們要先定義我們可能出現的字母及數字，總共有36個字。

alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'

另外用Image來寫一個開啟圖片的程式。

def img_loader(img_path):
    img = Image.open(img_path)
    return img.convert("RGB")

接下來我們要取出我們的圖片路徑跟圖片對應的label，我將註解跟程式碼寫在一起。

def make_dataset(data_path,ans_path,alphabet):
    img_names = os.listdir(data_path)# 取出圖片名稱
    img_names.sort(key=lambda x: int(x.split(".")[0]))# 讓圖片從小到大排序
    df_ans = pd.read_csv(ans_path)# 讀取label的CSV檔
    ans_list = list(df_ans["code"].values)# 取得label
    samples = []
    
    # 用zip將圖片跟對應的答案湊一對
    for ans, img_name in zip(ans_list, img_names):
        if len(str(ans)) == 5  :#num_char:
            
            # 將圖片名稱及路徑合併
            # 以便上述程式img_loader的執行
            img_path = os.path.join(data_path, img_name)
            target = []
            # 這邊做5個字的辨識，例如：A5GG2
            # 會轉換成target = [0,31,6,6,28]
            for char in str(ans):
                vec = [0] * 36 # num_class
                vec[alphabet.find(char)] = 1
                target += vec
                
            # 用samples把他們全部包起來    
            samples.append((img_path, target))
        else:
            print(img_name)
    return samples

接下來要把他從samples裡面一一讀取，並將label轉換成tensor。

class CaptchaData(Dataset):
    def __init__(self, data_path,ans_path,
                 transform=None, target_transform=None, alphabet=alphabet):
        super(Dataset, self).__init__()
        self.data_path = data_path
        self.ans_path = ans_path
        # self.num_class = num_class
        # self.num_char = num_char
        self.transform = transform
        self.target_transform = target_transform
        self.alphabet = alphabet
        self.samples = make_dataset(self.data_path,self.ans_path,self.alphabet
                                    )

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        img_path, target = self.samples[index]
        img = img_loader(img_path)
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, torch.Tensor(target)

下面這個程式是用來計算準確度的。

def calculat_acc(output, target):
    output, target = output.view(-1, 36), target.view(-1, 36)
    output = nn.functional.softmax(output, dim=1)
    output = torch.argmax(output, dim=1)
    target = torch.argmax(target, dim=1)
    output, target = output.view(-1, 5), target.view(-1, 5)
    correct_list = []
    for i, j in zip(target, output):
        if torch.equal(i, j):
            correct_list.append(1)
        else:
            correct_list.append(0)
    acc = sum(correct_list) / len(correct_list)
    return acc

接下來我們丟進model裡面train，我這邊選擇預訓練模型densenet201來做訓練。

batch_size = 15 # 依照個人設備去設定
base_lr = 0.6   # 設定優化器一開始的學習率，都可以嘗試看看
max_epoch = 60  # 看要練幾個epoch
model_path = './tset_densenet.pth'

def train():
    transforms = Compose([ToTensor()])
    train_dataset = CaptchaData('./pic_train2',
                                './answer/answer_train_v2.csv',
                                transform=transforms)
    train_data_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=0,
                                   shuffle=True, drop_last=True)
    test_data = CaptchaData('./pic_test2',
                            './answer/answer_test_v2.csv',
                            transform=transforms)
    test_data_loader = DataLoader(test_data, batch_size=batch_size,
                                  num_workers=0, shuffle=True, drop_last=True)
                                  
    # 使用densenet201來做訓練                         
    cnn = models.densenet201(num_classes=180)# 五個字,每個字有36種可能
    
    # 測試有沒有裝cuda，有沒有GPU可以使用
    if torch.cuda.is_available():
        cnn.cuda()
    if restor:
        cnn.load_state_dict(torch.load(model_path))
    # 這邊優化器我使用SGD+momentum，搭配CosineAnnealing（餘弦退火）的學習率scheduler，可以不固定學習率，以防他停在局部低點，有機會找到全局最佳解。
    optimizer = torch.optim.SGD(cnn.parameters(), lr=base_lr, momentum=0.9)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=16, eta_min=0, last_epoch=-1, verbose=False)
    criterion = nn.MultiLabelSoftMarginLoss()

    for epoch in range(max_epoch):
        start_ = time.time()

        loss_history = []
        acc_history = []
        cnn.train()

        for img, target in train_data_loader:
            img = Variable(img)
            target = Variable(target)
            if torch.cuda.is_available():
                img = img.cuda()
                target = target.cuda()
            output = cnn(img)
            loss = criterion(output, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            acc = calculat_acc(output, target)
            acc_history.append(float(acc))
            loss_history.append(float(loss))
        scheduler.step()
        print('train_loss: {:.4}|train_acc: {:.4}'.format(
            torch.mean(torch.Tensor(loss_history)),
            torch.mean(torch.Tensor(acc_history)),
        ))

        loss_history = []
        acc_history = []
        cnn.eval()
        for img, target in test_data_loader:
            img = Variable(img)
            target = Variable(target)
            if torch.cuda.is_available():
                img = img.cuda()
                target = target.cuda()
            output = cnn(img)

            acc = calculat_acc(output, target)
            acc_history.append(float(acc))
            loss_history.append(float(loss))
        print('test_loss: {:.4}|test_acc: {:.4}'.format(
            torch.mean(torch.Tensor(loss_history)),
            torch.mean(torch.Tensor(acc_history)),
        ))
        print('epoch: {}|time: {:.4f}'.format(epoch, time.time() - start_))
        print("========================================")
        torch.save(cnn.state_dict(), model_path)

練了6個epoch，準確度到快95%了，這裡就先用這個模型當作我們破解驗證碼的模型吧！！

• 今日小結

今天就把code分享給大家，明天來用今天練的模型測試看看可不可以直接進入網頁得到我們要的資訊吧！！！
明天見囉！！！